import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
from sklearn.metrics import confusion_matrix

def plot_confusion_matrix(df,threshold):
    def binarize(x, threshold):
        if x > threshold:
            x = 1
        else:
            x = 0
        return x

    df["binary"] = df["predictions"].apply(lambda x: binarize(x, threshold))
    true_labels = df["generated"].values
    predicted_labels = df["binary"].values

    classes = np.unique(np.concatenate((true_labels, predicted_labels)))
    cm = confusion_matrix(true_labels, predicted_labels, labels=classes)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm,
                annot=True,
                fmt="d",
                cmap="Blues",
                xticklabels=classes,
                yticklabels=classes)
    plt.xlabel("Predicted Labels")
    plt.ylabel("True Labels")
    plt.title("Confusion Matrix")
    plt.show()

def main():
    df = pd.read_csv("./output/3_test_multiple_results.csv")
    plot_confusion_matrix(df,0.5)

if __name__ == "__main__":
    main()